import pandas as pd
import numpy as np
import os
import glob
from tqdm import tqdm
import warnings
from datetime import datetime
import multiprocessing as mp # 导入多处理库
import math
import matplotlib.pyplot as plt # 导入绘图库
import seaborn as sns # 导入绘图库

# 忽略警告
warnings.filterwarnings('ignore')
# 设置绘图风格和处理负号显示
sns.set_theme(style="whitegrid")
plt.rcParams['axes.unicode_minus'] = False

# 定义MET值范围的常量
MET_SLEEP = 1.0  # MET值<1为睡眠
MET_STATIC = 1.6  # 1.0≤MET值<1.6之间为静态行为
MET_LOW = 3.0  # 1.6≤MET值<3.0为低强度
MET_MODERATE = 6.0  # 3.0≤MET值<6.0为中等强度
# MET值≥6.0为高强度

# 结果输出目录
RESULTS_DIR_TASK1 = "results_task1"
FIGURES_SAVE_PATH_TASK1 = os.path.join(RESULTS_DIR_TASK1, "figures")

# --- analyze_volunteer_data 函数保持不变，但为了并行化，需确保它是可序列化的且无副作用 ---
# (辅助函数：提取MET值和活动类型)
def extract_features_from_annotation(annotation):
    try:
        parts = str(annotation).split(';')
        activity_part = parts[0].strip()  # "7030 sleeping"
        activity_name = ' '.join(activity_part.split(' ')[1:])  # "sleeping"
        met_part = parts[1].strip()  # "MET 0.95"
        met_value = float(met_part.split(' ')[1])  # 0.95
        return activity_name, met_value
    except:
        return "unknown", 0.0

# (辅助函数：根据MET值分类活动)
def classify_activity_by_met(met_value):
    if met_value < MET_SLEEP:
        return 'sleep'
    elif met_value < MET_STATIC:
        return 'static'
    elif met_value < MET_LOW:
        return 'low'
    elif met_value < MET_MODERATE:
        return 'moderate'
    else:
        return 'high'

# (辅助函数：返回默认统计结果)
def get_default_stats(volunteer_id):
    return {
        '志愿者ID': volunteer_id,
        '记录总时长（小时）': 0,
        '睡眠总时长（分钟）': 0,
        '高等强度运动总时长（分钟）': 0,
        '中等强度运动总时长（分钟）': 0,
        '低等强度运动总时长（分钟）': 0,
        '静态活动总时长（分钟）': 0
    }

# --- 主要分析函数，将被并行调用 ---
def analyze_volunteer_data(file_path):
    """
    分析单个志愿者的数据 (设计为可在多进程中运行)
    
    参数:
        file_path (str): 数据文件路径
            
    返回:
        dict: 包含各项统计结果的字典
    """
    volunteer_id = os.path.basename(file_path).split('.')[0]
    
    try:
        try:
            df = pd.read_csv(file_path, low_memory=False)
        except Exception as e:
            # print(f"无法读取文件 {file_path}: {e}") # 在并行中打印可能混乱，可注释掉或使用日志
            return get_default_stats(volunteer_id)
            
        if 'time' not in df.columns or 'annotation' not in df.columns:
            # print(f"文件 {file_path} 缺少必要的列")
            return get_default_stats(volunteer_id)
            
        df = df.dropna(subset=['time', 'annotation'])
            
        if len(df) == 0:
            # print(f"文件 {file_path} 在清洗后没有有效数据")
            return get_default_stats(volunteer_id)
            
        # 处理时间戳
        first_time = df['time'].iloc[0]
        try:
            if isinstance(first_time, str) and '-' in first_time:
                df['datetime'] = pd.to_datetime(df['time'])
            else:
                df['time'] = pd.to_numeric(df['time'])
                df['datetime'] = pd.to_datetime(df['time'], unit='ms')
        except Exception as e:
            # print(f"处理时间戳时出错 {file_path}: {e}")
            try:
                df['datetime'] = pd.to_datetime(df['time'], errors='coerce')
                df = df.dropna(subset=['datetime'])
                if len(df) == 0:
                    # print(f"解析时间后，文件 {file_path} 没有有效数据")
                    return get_default_stats(volunteer_id)
            except:
                return get_default_stats(volunteer_id)
            
        df = df.sort_values('datetime')
        df['time_diff'] = df['datetime'].diff().dt.total_seconds()
            
        if len(df) > 1:
            median_interval = df['time_diff'].median()
            # 处理可能的NaN或非正数中位数
            if pd.isna(median_interval) or median_interval <= 0:
                median_interval = 0.1 # 假定一个小的默认间隔，如100ms
            df.loc[df.index[0], 'time_diff'] = median_interval
        else:
            median_interval = 0.1 # 只有一行数据时的默认间隔
            df.loc[df.index[0], 'time_diff'] = median_interval
            
        max_interval = 600  # 10分钟
        # 填充过大的时间差
        df.loc[df['time_diff'] > max_interval, 'time_diff'] = median_interval
        # 填充可能因diff产生的NaN或负值
        df['time_diff'] = df['time_diff'].fillna(median_interval).clip(lower=0)
            
        # 提取MET值
        try:
            df['activity_name'], df['met'] = zip(*df['annotation'].apply(extract_features_from_annotation))
        except Exception as e:
            # print(f"从标注提取特征时出错 {file_path}: {e}")
            df['met'] = df['annotation'].apply(lambda x: 
                                              float(str(x).split(';')[1].split(' ')[1]) 
                                              if ';' in str(x) and 'MET' in str(x) 
                                              else 0.0)
            
        # 计算总时长
        total_time_sec = df['time_diff'].sum()
        if total_time_sec <= 0 and len(df) > 1:
            total_time_sec = (df['datetime'].iloc[-1] - df['datetime'].iloc[0]).total_seconds()
        total_time_hours = total_time_sec / 3600
            
        # 分类活动
        df['activity_type'] = df['met'].apply(classify_activity_by_met)
            
        # 计算时长
        activity_minutes = {}
        for activity_type in ['sleep', 'static', 'low', 'moderate', 'high']:
            activity_time = df[df['activity_type'] == activity_type]['time_diff'].sum() / 60
            activity_minutes[activity_type] = activity_time
            
        return {
            '志愿者ID': volunteer_id,
            '记录总时长（小时）': total_time_hours,
            '睡眠总时长（分钟）': activity_minutes.get('sleep', 0),
            '高等强度运动总时长（分钟）': activity_minutes.get('high', 0),
            '中等强度运动总时长（分钟）': activity_minutes.get('moderate', 0),
            '低等强度运动总时长（分钟）': activity_minutes.get('low', 0),
            '静态活动总时长（分钟）': activity_minutes.get('static', 0)
        }
            
    except Exception as e:
        # print(f"处理文件 {file_path} 时出现意外错误: {e}") # 避免在并行时过多打印
        return get_default_stats(volunteer_id)

# --- 并行分析所有志愿者 --- #
def analyze_all_volunteers_parallel(data_dir, num_processes=None):
    """
    并行分析所有志愿者的数据并生成结果
    
    参数:
        data_dir (str): 数据文件夹路径
        num_processes (int, optional): 使用的进程数。默认为CPU核心数。
        
    返回:
        DataFrame: 包含所有志愿者统计结果的数据框
    """
    files = glob.glob(os.path.join(data_dir, "P*.csv"))
    
    if not files:
        print(f"在目录 {data_dir} 中未找到任何P*.csv文件")
        return pd.DataFrame()
        
    print(f"找到 {len(files)} 个志愿者数据文件，将使用并行处理...")
    
    # 确定进程数
    if num_processes is None:
        num_processes = mp.cpu_count()
    print(f"使用 {num_processes} 个进程进行并行计算")
    
    results = []
    # 创建进程池
    with mp.Pool(processes=num_processes) as pool:
        # 使用imap_unordered和tqdm显示进度
        # imap_unordered相比map效率稍高，因为它不保证返回顺序
        # 使用tqdm包装pool.imap_unordered来显示进度条
        with tqdm(total=len(files), desc="并行处理志愿者数据") as pbar:
            for result in pool.imap_unordered(analyze_volunteer_data, files):
                results.append(result)
                pbar.update(1) # 每完成一个任务，进度条更新一次
    
    # 将结果转换为DataFrame
    if not results:
        print("处理结果为空")
        return pd.DataFrame()
        
    result_df = pd.DataFrame(results)
    
    # 调整列的顺序
    columns = [
        '志愿者ID', 
        '记录总时长（小时）', 
        '睡眠总时长（分钟）', 
        '高等强度运动总时长（分钟）', 
        '中等强度运动总时长（分钟）', 
        '低等强度运动总时长（分钟）', 
        '静态活动总时长（分钟）'
    ]
        
    # 确保所有列都存在
    for col in columns:
        if col not in result_df.columns:
            result_df[col] = 0
        
    result_df = result_df[columns]
        
    # 对结果进行排序
    result_df = result_df.sort_values('志愿者ID')
        
    # 设置小数点后4位
    for col in columns[1:]:
        result_df[col] = result_df[col].round(4)
        
    return result_df

# --- 新增绘图函数 (使用英文标签) --- #

def plot_activity_distribution_stacked_bar(df, save_path):
    """绘制每个志愿者的活动时长分布堆叠条形图"""
    os.makedirs(save_path, exist_ok=True)
    if df.empty:
        print("Info: No data to plot stacked bar chart.")
        return

    plot_df = df.set_index('志愿者ID')
    activity_cols = [
        '睡眠总时长（分钟）', '静态活动总时长（分钟）', '低等强度运动总时长（分钟）',
        '中等强度运动总时长（分钟）', '高等强度运动总时长（分钟）'
    ]
    # 计算占比
    plot_df_perc = plot_df[activity_cols].apply(lambda x: x / x.sum() * 100, axis=1)
    # 重命名列为英文以便图例显示
    plot_df_perc.columns = ['Sleep', 'Static', 'Low', 'Moderate', 'High']

    plt.figure(figsize=(15, max(8, len(plot_df_perc) * 0.3))) # 动态调整高度
    plot_df_perc.plot(kind='barh', stacked=True, colormap='viridis')

    plt.title('Activity Duration Distribution per Subject (%)')
    plt.xlabel('Percentage of Total Time (%)')
    plt.ylabel('Subject ID')
    plt.legend(title='Activity Type', bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.tight_layout(rect=[0, 0, 0.85, 1]) # 调整布局给图例留空间
    plt.savefig(os.path.join(save_path, "activity_distribution_stacked_bar.png"), dpi=300)
    plt.close()

def plot_average_activity_pie(df, save_path):
    """绘制所有志愿者平均活动时长占比饼图"""
    os.makedirs(save_path, exist_ok=True)
    if df.empty:
        print("Info: No data to plot average activity pie chart.")
        return

    activity_cols = [
        '睡眠总时长（分钟）', '静态活动总时长（分钟）', '低等强度运动总时长（分钟）',
        '中等强度运动总时长（分钟）', '高等强度运动总时长（分钟）'
    ]
    # 计算每个类别的总时长（所有志愿者）
    total_durations = df[activity_cols].sum()
    # 重命名索引为英文
    total_durations.index = ['Sleep', 'Static', 'Low', 'Moderate', 'High']
    # 过滤掉时长为0的类别，避免饼图错误
    total_durations = total_durations[total_durations > 0]

    if total_durations.empty:
        print("Info: All activity durations are zero, cannot plot pie chart.")
        return

    plt.figure(figsize=(8, 8))
    plt.pie(total_durations, labels=total_durations.index, autopct='%1.1f%%', startangle=90, colors=sns.color_palette("viridis", len(total_durations)))
    plt.title('Average Activity Duration Distribution (%)')
    plt.savefig(os.path.join(save_path, "average_activity_pie_chart.png"), dpi=300)
    plt.close()

def plot_total_duration_histogram(df, save_path):
    """绘制总记录时长分布直方图"""
    os.makedirs(save_path, exist_ok=True)
    if df.empty or '记录总时长（小时）' not in df.columns:
        print("Info: No data or missing column for total duration histogram.")
        return

    plt.figure(figsize=(10, 6))
    sns.histplot(df['记录总时长（小时）'], kde=True, bins=20)
    plt.title('Distribution of Total Recording Duration per Subject')
    plt.xlabel('Total Recording Duration (hours)')
    plt.ylabel('Number of Subjects')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.savefig(os.path.join(save_path, "total_duration_histogram.png"), dpi=300)
    plt.close()

def plot_activity_duration_boxplots(df, save_path):
    """绘制各活动类型时长箱线图"""
    os.makedirs(save_path, exist_ok=True)
    if df.empty:
        print("Info: No data to plot activity duration boxplots.")
        return

    activity_cols_minutes = {
        'Sleep': '睡眠总时长（分钟）',
        'Static': '静态活动总时长（分钟）',
        'Low Intensity': '低等强度运动总时长（分钟）',
        'Moderate Intensity': '中等强度运动总时长（分钟）',
        'High Intensity': '高等强度运动总时长（分钟）'
    }
    # 将数据从宽格式转换为长格式，便于绘图
    df_long = df.melt(id_vars=['志愿者ID'], value_vars=list(activity_cols_minutes.values()),
                      var_name='Activity Type Original', value_name='Duration (minutes)')

    # 创建映射以转换为英文标签
    eng_labels_map = {v: k for k, v in activity_cols_minutes.items()}
    df_long['Activity Type'] = df_long['Activity Type Original'].map(eng_labels_map)

    plt.figure(figsize=(12, 7))
    sns.boxplot(x='Activity Type', y='Duration (minutes)', data=df_long, palette='viridis')
    plt.title('Distribution of Durations for Each Activity Type')
    plt.xlabel('Activity Type')
    plt.ylabel('Duration (minutes)')
    plt.xticks(rotation=15) # 轻微旋转标签防止重叠
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "activity_duration_boxplots.png"), dpi=300)
    plt.close()

def main():
    """主函数"""
    # 创建结果目录
    os.makedirs(RESULTS_DIR_TASK1, exist_ok=True)
    os.makedirs(FIGURES_SAVE_PATH_TASK1, exist_ok=True)

    # 设置数据目录 (请确保这是服务器上的正确路径)
    data_dir = r"/share/home/u2201230214/GPR/B/附件1"
    
    if not os.path.exists(data_dir):
        print(f"路径 {data_dir} 不存在，请检查路径设置")
        return
    
    # 进行并行分析
    print("开始并行分析志愿者活动数据...")
    result_df = analyze_all_volunteers_parallel(data_dir)
    
    if result_df.empty:
        print("Analysis result is empty, check data path and files.") # 改为英文
        return
    
    # 保存结果
    output_file = os.path.join(RESULTS_DIR_TASK1, "result_1_parallel.xlsx") # 保存到结果目录
    try:
        # 需要安装 openpyxl: pip install openpyxl
        result_df.to_excel(output_file, index=False)
        print(f"Analysis complete. Results saved to {output_file}") # 改为英文
    except Exception as e:
        print(f"Error saving Excel results: {e}") # 改为英文
        try:
            csv_file = os.path.join(RESULTS_DIR_TASK1, "result_1_parallel.csv") # 保存到结果目录
            result_df.to_csv(csv_file, index=False)
            print(f"Results saved as CSV: {csv_file}") # 改为英文
        except Exception as csv_err:
            print(f"Failed to save as CSV: {csv_err}") # 改为英文
    
    # --- 调用绘图函数 --- #
    print("\nGenerating plots...")
    try:
        plot_activity_distribution_stacked_bar(result_df, FIGURES_SAVE_PATH_TASK1)
        plot_average_activity_pie(result_df, FIGURES_SAVE_PATH_TASK1)
        plot_total_duration_histogram(result_df, FIGURES_SAVE_PATH_TASK1)
        plot_activity_duration_boxplots(result_df, FIGURES_SAVE_PATH_TASK1)
        print(f"Plots saved to {FIGURES_SAVE_PATH_TASK1}")
    except Exception as plot_err:
        print(f"Error generating plots: {plot_err}")

    # 显示部分结果
    print("\nResults Preview:") # 改为英文
    print(result_df.head())

if __name__ == "__main__":
    # Windows平台需要此保护，Linux/macOS也建议加上
    mp.freeze_support() # 对于打包成可执行文件尤其重要
    main() 